import numpy as np
from Tools import metrics_result
from Tools import readbunchobj
from sklearn.ensemble import RandomForestClassifier
from joblib import dump
from sklearn.preprocessing import StandardScaler

# 导入训练集
trainpath = "../train_word_bag/tf_idf_space.dat"
train_set = readbunchobj(trainpath)

# 导入测试集
testpath = "../test_word_bag/test_tf_idf_space.dat"
test_set = readbunchobj(testpath)

# 输出单词矩阵的类型
print("标准化/归一化前的矩阵形状:")
print(np.shape(train_set.tdm))
print(np.shape(test_set.tdm))

# 对数据进行标准化或归一化
scaler = StandardScaler(with_mean=False)
train_set.tdm = scaler.fit_transform(train_set.tdm)
test_set.tdm = scaler.transform(test_set.tdm)

# 随机森林
print("开始训练随机森林模型...")
# 一个随机森林分类器，它包含100棵决策树，使用所有可用的CPU核心进行并行计算，并且设置了一个固定的随机数生成器种子以确保模型的可复现性。
rf_clf = RandomForestClassifier(n_estimators=100, n_jobs=-1, random_state=0)
rf_clf.fit(train_set.tdm, train_set.label)
print("随机森林模型训练完成。")
dump(rf_clf, 'RF_model.joblib')

# 预测分类结果
print("USE: 随机森林模型 预测分类结果")
rf_predicted = rf_clf.predict(test_set.tdm)
rf_total = len(rf_predicted)
print("结束")

# 性能评估
print("\n随机森林模型：")
metrics_result(test_set.label, rf_predicted, rf_total)
